import matplotlib.pyplot as plt
import numpy as np
import os
from config import *

def visualize_grid(grid, title="Environment", save_path=None):
    """可视化环境网格"""
    colors = {
        EMPTY: 'white',
        START: 'green',
        AGENT: 'blue',
        COIN: 'yellow',
        OBSTACLE: 'black',
        END: 'red'
    }
    
    names = {
        EMPTY: 'Empty',
        START: 'Start',
        AGENT: 'Agent',
        COIN: 'Coin',
        OBSTACLE: 'Obstacle',
        END: 'End'
    }
    
    # 创建网格可视化
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_title(title)
    
    # 绘制网格
    height, width = grid.shape
    for i in range(height):
        for j in range(width):
            cell_value = grid[i, j]
            color = colors.get(cell_value, 'gray')  # 使用灰色作为未知值
            rect = plt.Rectangle((j, height - i - 1), 1, 1, color=color, ec='black')
            ax.add_patch(rect)
            
            # 添加文本标签，除了空白和障碍物
            if cell_value not in [EMPTY, OBSTACLE]:
                ax.text(j + 0.5, height - i - 0.5, names.get(cell_value, 'Unknown'),
                        ha='center', va='center', fontsize=8)
    
    # 设置坐标轴范围和标签
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    ax.set_xticks(range(width))
    ax.set_yticks(range(height))
    ax.set_xticklabels(range(width))
    ax.set_yticklabels(range(height-1, -1, -1))
    ax.grid(True)
    
    # 保存或显示图像
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()

def visualize_prediction(true_state, pred_state, action, save_path=None):
    """可视化真实状态和预测状态的比较"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))
    
    # 绘制真实状态
    visualize_state(true_state, ax1, title=f"True State After Action: {ACTION_NAMES[action]}")
    
    # 确保预测状态中有智能体
    if not np.any(pred_state == AGENT):
        # 如果预测中没有智能体，使用真实状态中的智能体位置
        agent_pos = np.where(true_state == AGENT)
        if len(agent_pos[0]) > 0:
            # 复制预测状态以避免修改原始数据
            pred_state_copy = pred_state.copy()
            # 放置智能体
            pred_state_copy[agent_pos[0][0], agent_pos[1][0]] = AGENT
            # 绘制修改后的预测状态
            visualize_state(pred_state_copy, ax2, title="Predicted State (Agent position added)")
        else:
            # 如果真实状态也没有智能体，正常绘制
            visualize_state(pred_state, ax2, title="Predicted State")
    else:
        # 正常绘制预测状态
        visualize_state(pred_state, ax2, title="Predicted State")
    
    # 保存或显示图像
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()

def visualize_state(grid, ax, title="Environment"):
    """在给定的轴上可视化状态"""
    colors = {
        EMPTY: 'white',
        START: 'green',
        AGENT: 'blue',
        COIN: 'yellow',
        OBSTACLE: 'black',
        END: 'red'
    }
    
    names = {
        EMPTY: 'Empty',
        START: 'Start',
        AGENT: 'Agent',
        COIN: 'Coin',
        OBSTACLE: 'Obstacle',
        END: 'End'
    }
    
    ax.set_title(title)
    
    # 绘制网格
    height, width = grid.shape
    for i in range(height):
        for j in range(width):
            cell_value = grid[i, j]
            color = colors.get(cell_value, 'gray')  # 使用灰色作为未知值
            rect = plt.Rectangle((j, height - i - 1), 1, 1, color=color, ec='black')
            ax.add_patch(rect)
            
            # 添加文本标签，除了空白和障碍物
            if cell_value not in [EMPTY, OBSTACLE]:
                ax.text(j + 0.5, height - i - 0.5, names.get(cell_value, 'Unknown'),
                        ha='center', va='center', fontsize=8)
    
    # 设置坐标轴范围和标签
    ax.set_xlim(0, width)
    ax.set_ylim(0, height)
    ax.set_xticks(range(width))
    ax.set_yticks(range(height))
    ax.set_xticklabels(range(width))
    ax.set_yticklabels(range(height-1, -1, -1))
    ax.grid(True)

def plot_training_progress(train_losses, val_losses, 
                           train_env_accuracies, val_env_accuracies,
                           train_agent_pos_accuracies=None, val_agent_pos_accuracies=None,
                           train_agent_mov_accuracies=None, val_agent_mov_accuracies=None,
                           save_path="training_progress.png"):
    """绘制增强的训练进度图表"""
    # 确定图表数量
    num_plots = 3 if train_agent_mov_accuracies is not None else 2
    
    # 创建图表
    fig, axes = plt.subplots(num_plots, 1, figsize=(15, 5*num_plots))
    
    # 绘制损失
    axes[0].plot(train_losses, label='Train Loss')
    axes[0].plot(val_losses, label='Validation Loss')
    axes[0].set_title('Total Loss')
    axes[0].set_xlabel('Epoch')
    axes[0].set_ylabel('Loss')
    axes[0].legend()
    axes[0].grid(True)
    
    # 绘制环境准确率
    axes[1].plot(train_env_accuracies, label='Train Env Accuracy')
    axes[1].plot(val_env_accuracies, label='Validation Env Accuracy')
    axes[1].set_title('Environment State Accuracy')
    axes[1].set_xlabel('Epoch')
    axes[1].set_ylabel('Accuracy')
    axes[1].set_ylim(0, 1.05)
    axes[1].legend()
    axes[1].grid(True)
    
    # 绘制智能体位置和移动准确率
    if num_plots > 2:
        # 位置准确率
        axes[2].plot(train_agent_pos_accuracies, label='Train Agent Position Accuracy')
        axes[2].plot(val_agent_pos_accuracies, label='Validation Agent Position Accuracy')
        
        # 移动准确率
        if train_agent_mov_accuracies is not None:
            axes[2].plot(train_agent_mov_accuracies, '--', label='Train Agent Movement Accuracy')
            axes[2].plot(val_agent_mov_accuracies, '--', label='Validation Agent Movement Accuracy')
        
        axes[2].set_title('Agent Position & Movement Accuracy')
        axes[2].set_xlabel('Epoch')
        axes[2].set_ylabel('Accuracy')
        axes[2].set_ylim(0, 1.05)
        axes[2].legend()
        axes[2].grid(True)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图像
    plt.savefig(save_path)
    plt.close()